import java.util.TreeMap;

// 含有嵌套的分子式求原子数量
// 测试链接 : https://leetcode.cn/problems/number-of-atoms/
public class NumberOfAtoms {
    public int where;

    public String countOfAtoms(String formula) {
        where = 0;
        TreeMap<String, Integer> map = new TreeMap<>();
        map = func(formula.toCharArray(), 0);

        StringBuilder ans = new StringBuilder();
        for(String key : map.keySet() ){
            ans.append(key);
            int cnt = map.get(key);

            if(cnt > 1){
                ans.append(cnt);
            }
        }

        return ans.toString();
    }


    public TreeMap<String, Integer> func(char[] s, int i){
        TreeMap<String, Integer> ans = new TreeMap<>();

        int cnt = 0; //记录个数
        StringBuilder name = new StringBuilder(); //记录历史名字
        TreeMap<String, Integer> pre = null;

        while(i < s.length && s[i] != ')'){
            if(s[i] >= 'A' && s[i] <= 'Z' || s[i] == '('){
                fill(ans, name, pre, cnt); //结算历史
                //清空历史
                name.setLength(0);
                pre = null;
                cnt = 0;

                if(s[i] >= 'A' && s[i] <= 'Z'){
                    name.append(s[i++]);
                }else{
                    pre = func(s, i+1);
                    i = where + 1;
                }
            }
            else if(s[i] >= 'a' && s[i] <= 'z'){
                name.append(s[i++]);
            }
            else{
                cnt = cnt * 10 + s[i++] - '0';
            }
        }

        fill(ans, name, pre, cnt);
        where = i;
        return ans;
    }

    public void fill(TreeMap<String, Integer> ans, StringBuilder name, TreeMap<String, Integer> pre, int cnt) {
        if(name.length() > 0 || pre != null){
            cnt = cnt == 0 ? 1 : cnt;
            if(name.length() > 0){
                String key = name.toString();
                ans.put(key, ans.getOrDefault(key,0) + cnt);
            }
            else {
                for(String key : pre.keySet()){
                    int num = pre.get(key) * cnt;
                    ans.put(key, ans.getOrDefault(key,0) + num);
                }
            }
        }
    }
}
